/*  QuickLSB is steganography program for hiding files
    Copyright (C) 2010  Aleksandrs Himulya and Anton Tarasov
 
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License.
 
    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.
 
    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>. */

#include <stdio.h>
#include <cstdlib>


using namespace std;


// òèï, îïèñûâàþùèé áàéòîâûé ìàññèâ
typedef struct bytearray_struct{
        char* array;
        long size;
} bytearray;

// ôóíêöèÿ ïðèíèìàåò èìÿ ôàéëà è âûäàåò åãî ñîäåðæèìîå â âèäå áàéòîâîãî ìàññèâà
bytearray* loadFile(char* file_path)
{
        FILE* f = fopen(file_path, "rb");
         
        if ( f == NULL )
                return NULL;

        // óçíàåì ðàçìåð
        fseek(f, 0, SEEK_END);
        long size = ftell(f);
        fseek(f, 0, SEEK_SET);

        // ñîçäàåì áàéòîâûé ìàññèâ
        bytearray* ba = new bytearray;
        ba->array = new char[size];
        ba->size = size;

        // ñ÷èòûâàåì äàííûå èç ôàéëà è çàêðûâàåì åãî
        fread(ba->array, size, 1, f);
        fclose(f);

        // âîçâðàùàåì áàéòîâûé ìàññèâ
        return ba;
}

// ñîõðàíÿåò áàéòîâûé ìàññèâ â ôàéë
int saveToFile(char* file_path, bytearray* ba)
{
        FILE* f = fopen(file_path, "wb+");
        if ( f == NULL )
                return 1;
        fwrite(ba->array, ba->size, 1, f);
        fclose(f);
        return 0;
}

bool getBit(char byte, char n)
{
    // let's count bits in this way: 87654321
    char mask = 1;
    mask = mask << (n -1);
    char result = mask & byte;
    return bool(result);
}


char setBit(char byte, char n, bool bit)
{
    char mask = 1;
    mask = mask << (n - 1);
    byte = byte & ~mask;
    byte = byte | (mask & (int(bit) << (n - 1)));
    return byte;
}

// óïàêîâêà äàííûõ
bytearray* packFile(bytearray* pic, bytearray* secret, int n)
{
        int pic_byte_i = 1024; // íîìåð àêòèâíîãî áàéòà èçîáðàæåíèÿ
        int pic_bit_i = 1; // íîìåð àêòèâíîãî áèòà èçîáðàæåíèÿ
        
        // çàïèñûâàåì êîëè÷åñòâî èñïîëüçóåìûõ áèò      
        int a[3];        
        switch (n)
        {
                case 1: a[0] = 0; a[1] = 0; a[2] = 0; break;
                case 2: a[0] = 0; a[1] = 0; a[2] = 1; break;
                case 3: a[0] = 0; a[1] = 1; a[2] = 0; break;
                case 4: a[0] = 0; a[1] = 1; a[2] = 1; break;
                case 5: a[0] = 1; a[1] = 0; a[2] = 0; break;
                case 6: a[0] = 1; a[1] = 0; a[2] = 1; break;
                case 7: a[0] = 1; a[1] = 1; a[2] = 0; break;
                case 8: a[0] = 1; a[1] = 1; a[2] = 1; break;
        }
        pic->array[ pic_byte_i ] = setBit( pic->array[ pic_byte_i ], 1, a[0]);
        pic_byte_i++;
        pic->array[ pic_byte_i ] = setBit( pic->array[ pic_byte_i ], 1, a[1]);
        pic_byte_i++;
        pic->array[ pic_byte_i ] = setBit( pic->array[ pic_byte_i ], 1, a[2]);
        pic_byte_i++;

        // çàïèñûâàåì êîëè÷åñòâî áàéò
        int size = secret->size;
        int byte_index = 0;
        int bit_index = 1;
        while ( byte_index < sizeof(int) )
        {
                bool bit = getBit(  ((char*)&size)[byte_index], bit_index );
                pic->array[pic_byte_i] = setBit( pic->array[pic_byte_i], pic_bit_i, bit );
                bit_index++;
                pic_bit_i++;
                if (pic_bit_i > n)
                {
                        pic_bit_i = 1;
                        pic_byte_i++;
                }
                if (bit_index > 8)
                {
                        bit_index = 1;
                        byte_index++;
                }
        }

        // íà÷èíàåì ïàêîâàòü áàéòû ñåêðåòíîãî ôàéëà
        int secret_byte_i = 0; // íîìåð àêòèâíîãî áàéòà èçîáðàæåíèÿ
        int secret_bit_i = 1; // íîìåð àêòèâíîãî áèòà èçîáðàæåíèÿ
        while ( secret_byte_i < size )
        {
                bool bit = getBit( secret->array[secret_byte_i], secret_bit_i );
                pic->array[pic_byte_i] = setBit( pic->array[pic_byte_i], pic_bit_i, bit );
                secret_bit_i++;
                pic_bit_i++;
                if (pic_bit_i > n)
                {
                        pic_bit_i = 1;
                        pic_byte_i++;
                }
                if (secret_bit_i > 8)
                {
                        secret_bit_i = 1;
                        secret_byte_i++;
                }
                if (pic_byte_i > pic->size)
                {
                        printf("Error: secret file is too large!");
                        return NULL;
                }
        }

        // ïèøåì ìóñîð
        while ( pic_byte_i < pic->size )
        {
                pic->array[pic_byte_i] = setBit( pic->array[pic_byte_i], pic_bit_i, bool(rand()%2) );
                pic_bit_i++;
                if (pic_bit_i > n)
                {
                        pic_bit_i = 1;
                        pic_byte_i++;
                }
        }  
        return pic; 
}



// ðàñïàêîâêà äàííûõ
bytearray* unpackFile(bytearray* pic)
{
        int pic_byte_i = 1024; // íîìåð àêòèâíîãî áàéòà èçîáðàæåíèÿ
        int pic_bit_i = 1; // íîìåð àêòèâíîãî áèòà èçîáðàæåíèÿ

        bool a[3];
        a[0] = getBit(pic->array[pic_byte_i], 1);
        pic_byte_i++;
        a[1] = getBit(pic->array[pic_byte_i], 1);
        pic_byte_i++;
        a[2] = getBit(pic->array[pic_byte_i], 1);
        pic_byte_i++;
        int mask = int(a[0])*100 + int(a[1])*10 + int(a[2]);
        int n =  0; // êîëè÷åñòâî èñïîëüçóåìûõ áàéò
        switch (mask)
        {
                case 0: n = 1; break;
                case 1: n = 2; break;
                case 10: n = 3; break;
                case 11: n = 4; break;
                case 100: n = 5; break;
                case 101: n = 6; break;
                case 110: n = 7; break;
                case 111: n = 8; break;
        }

        // ïîëó÷àåì êîëè÷åñòâî áàéò â ñåêðåòíîì ñîîáùåíèè
        char size_array[4];

        int byte_index = 0;
        int bit_index = 1;
        while ( byte_index < sizeof(int) )
        {
                bool bit = getBit( pic->array[pic_byte_i], pic_bit_i );  
                 size_array[byte_index] = setBit( size_array[byte_index], bit_index, bit );;
                bit_index++;
                pic_bit_i++;
                if (pic_bit_i > n)
                {
                        pic_bit_i = 1;
                        pic_byte_i++;
                }
                if (bit_index > 8)
                {
                        bit_index = 1;
                        byte_index++;
                }
        }
        int size = *((int*)size_array);

       
        // íà÷èíàåì ðàñïàêîâûâàòü áàéòû ñåêðåòíîãî ôàéëà
        bytearray* ba = new bytearray;
        ba->size = size;
        ba->array = new char[size];
        int secret_byte_i = 0; // íîìåð àêòèâíîãî áàéòà ñåêðåòíîãî ôàéëà
        int secret_bit_i = 1; // íîìåð àêòèâíîãî áèòà ñåêðåòíîãî ôàéëà
        while ( secret_byte_i < size )
        {
                bool bit = getBit( pic->array[pic_byte_i], pic_bit_i );
                ba->array[secret_byte_i] = setBit( ba->array[secret_byte_i], secret_bit_i, bit );
                secret_bit_i++;
                pic_bit_i++;
                if (pic_bit_i > n)
                {
                        pic_bit_i = 1;
                        pic_byte_i++;
                }
                if (secret_bit_i > 8)
                {
                        secret_bit_i = 1;
                        secret_byte_i++;
                }
        }
        return ba;
}

int enter_num_bit(long unsigned int n_pic)//задание номера используемых бит в байте
{
    char input;
    int num;
    printf("Введите номер бита, который и биты после которого будут заменятся.\n При выборе значения 8 возможно будет скрыть до %d байт информации и при этом изменеия в качестве изображеия практически не произойдет.\n При выборе значения 7 возможно будет скрыть до %d байт информации и изменения будут малозаметны.\n При выборе значения 6 возможно будет скрыть до %d байт, изменения будут заметны.\n При выборе значения 5 возможно будет скрыть до %d байт, но изменения будут сильно заметны.\n При выборе значения 4 возможно будет скрыть до %d байт, изображение сильно изменится.\n При выборе значения 3 возможно будет скрыть до %d байт, изображение скорее всего будет пратически неразличимо.\n При выборе значения 2 возможно будет скрыть до %d байт, исходное изображение будет не различить.\n При выборе значения 1 возможно будет скрыть до %d байт, вы просто скопируете файл и не увидите на изображении или его части ничего кроме шумов.\n Ваш выбор [1/2/3/4/5/6/7/8]:", ((n_pic-1064)/8), ((n_pic-1064)/8*2), ((n_pic-1064)/8*3), ((n_pic-1064)/8*4), ((n_pic-1064)/8*5), ((n_pic-1064)/8*6), ((n_pic-1064)/8*7), (n_pic-1064));
    scanf("%1d", &input);
    if (input == 8)
        num = 1;
    else if (input == 7)
        num = 2;
    else if (input == 6)
        num = 3;
    else if (input == 5)
        num = 4;
    else if (input == 4)
        num = 5;
    else if (input == 3)
        num = 6;
    else if (input == 2)
        num = 7;
    else if (input == 1)
        num = 8;
    else
        printf("Wrong number\n");
    return num;
}



void prepare_pack(void)
{
    char path_file_sec[255], path_file_open[255], path_file_stg[255];
    printf("Введите полный путь к открытому файлу:");
    scanf("%s", &path_file_open);
    bytearray* pic = loadFile((char*)&path_file_open);//открытие файла изображения
    int n = enter_num_bit(pic->size);//задание номера бита, который будет меняться
    printf("Введите полный путь к скрываемому файлу:");
    scanf("%s", &path_file_sec);
    bytearray* secret = loadFile((char*)&path_file_sec);//открытие файла со скрытой информацией
    if (pic == NULL || secret == NULL) //проверка, открыты ли файлы
    {
        printf("Возникла ошибка при чтении файла(ов)\n");
        return;
    }
    else
        printf("Файлы успешно открыты\n");
    long p_size = pic->size;
    long s_size = secret->size;
    if (((p_size)-1064)<((s_size)*8*n))
    {
        printf("Открытый файл слишком мал для скрываемого или скрываемый файл слишком большой для открытого. Попробуйте увеличить номер бита, после которого записывается информация или заменить открытый и/или скрытый файл(ы).\n");
        return;
    }
    printf("Введите полный путь к файлу, в который будет записан получившийся стегофайл:");
    scanf("%s", &path_file_stg);
    bytearray* stg = packFile(pic, secret, n);
    int err = saveToFile((char*)&path_file_stg, stg);
    if (err != 0)
        printf("Возникла ошибка при записи\n");
    else
    {};
    return;
}

void prepare_unpack(void)
{
    char path_file_stg[255], path_ext_file[255];
    printf("Введите полный путь к стегофайлу:");
    scanf("%s", &path_file_stg);
    bytearray* stg = loadFile((char*)&path_file_stg);//открытие файла изображения со скрытыми данными
    printf("Введите путь к файлу, в который будет распакован скрытый файл:");
    scanf("%s", &path_ext_file);
    if (stg == NULL) //проверка, открыты ли файлы
    {
        printf("Возникла ошибка при чтении файла\n");
        return;
    }
    else
        printf("Файл успешно открыт\n");
    bytearray* secret = unpackFile(stg);
    int err = saveToFile((char*)&path_ext_file, secret);
    if (err != 0)
        printf("Возникла ошибка при записи\n");
    else
    {};
    return;
}

int main()
{
    char com;
    printf("Эта программа для создания стегофайлов[c(create/создать)/e(extract/распаковать)/q(quit/выйти)]:");
    scanf("%1c", &com);
    switch(com)//переход к...
    {
        case 'c': prepare_pack();//созданию стего-файла
        case 'e': prepare_unpack();//извлечению информации из стего-файла
        case 'q': return 0;//выход
        default:  printf ("Wrong command\n");//ошибка в случае неверной команды
    }
    return 0;
}
